{
 "cells": [
  {
   "metadata": {},
   "cell_type": "raw",
   "source": "因果模型训练",
   "id": "aeb32d87ceb7e5ba"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-05-21T14:23:07.054828Z",
     "start_time": "2025-05-21T14:23:01.604805Z"
    }
   },
   "source": [
    "from datasets import load_dataset\n",
    "from MyHelper import *\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling, TrainingArguments, Trainer, DataCollatorForWholeWordMask\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ],
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\51165\\.conda\\envs\\e12\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "加载数据集",
   "id": "3926ff3a848572dd"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-21T14:23:09.195261Z",
     "start_time": "2025-05-21T14:23:07.057832Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_name = \"Langboat/bloom-389m-zh\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)"
   ],
   "id": "8454a1d4f5d6abb1",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-21T14:23:14.147527Z",
     "start_time": "2025-05-21T14:23:09.316837Z"
    }
   },
   "cell_type": "code",
   "source": [
    "dataset_name = \"pleisto/wikipedia-cn-20230720-filtered\"\n",
    "ds = load_dataset(dataset_name)\n",
    "ds, ds.get(\"train\")[0]"
   ],
   "id": "a26b1c503b94eab9",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DatasetDict({\n",
       "     train: Dataset({\n",
       "         features: ['completion', 'source'],\n",
       "         num_rows: 254547\n",
       "     })\n",
       " }),\n",
       " {'completion': '昭通机场（ZPZT）是位于中国云南昭通的民用机场，始建于1935年，1960年3月开通往返航班“昆明－昭通”，原来属军民合用机场。1986年机场停止使用。1991年11月扩建，于1994年2月恢复通航。是西南地区「文明机场」，通航城市昆明。 机场占地1957亩，飞行区等级为4C，有一条跑道，长2720米，宽48米，可供波音737及以下机型起降。机坪面积6600平方米，停机位2个，航站楼面积1900平方米。位于城东6公里处，民航路与金鹰大道交叉处。\\n航点\\n客服电话\\n昭通机场客服电话：0870-2830004',\n",
       "  'source': 'wikipedia.zh2307'})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-21T14:23:14.770644Z",
     "start_time": "2025-05-21T14:23:14.171728Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def process_function(examples, tokenizer=tokenizer):\n",
    "    return tokenizer(examples[\"completion\"], max_length=384, truncation=True)\n",
    "ds = ds.map(process_function, batched=True, num_proc=16, batch_size=64, remove_columns=ds[\"train\"].column_names)"
   ],
   "id": "af29533ff3cd41b3",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-21T14:24:08.739854Z",
     "start_time": "2025-05-21T14:23:14.779982Z"
    }
   },
   "cell_type": "code",
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"output/0207/\",\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=16,\n",
    "    max_steps=30,\n",
    "    logging_steps=10\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "\n",
    "trainer.train()"
   ],
   "id": "5947e3fff3cf1ca7",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='30' max='30' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [30/30 00:51, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>4.031000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>4.040700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>3.869600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=30, training_loss=3.980426788330078, metrics={'train_runtime': 53.1521, 'train_samples_per_second': 9.031, 'train_steps_per_second': 0.564, 'total_flos': 238935266476032.0, 'train_loss': 3.980426788330078, 'epoch': 0.001885702836804205})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-21T14:24:08.758066Z",
     "start_time": "2025-05-21T14:24:08.755931Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "cccc5931305e5635",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
